package com.ajocer.springbootinit.monitor;

import com.ajocer.springbootinit.config.MilvusConfig;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.R;
import io.milvus.param.collection.HasCollectionParam;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Counter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 服务健康检查和自动恢复机制
 */
@Component
@EnableScheduling
@Slf4j
public class ServiceHealthChecker {

    @Autowired
    private MilvusServiceClient milvusClient;

    @Autowired
    private MilvusConfig milvusConfig;

    @Autowired
    private RedisConnectionFactory redisConnectionFactory;

    @Autowired
    private MeterRegistry meterRegistry;

    // 服务状态指标
    private Counter milvusFailureCounter;
    private Counter milvusRecoveryCounter;
    private Counter redisFailureCounter;
    private Counter redisRecoveryCounter;

    // 服务状态标志
    private final AtomicBoolean milvusHealthy = new AtomicBoolean(true);
    private final AtomicBoolean redisHealthy = new AtomicBoolean(true);

    @PostConstruct
    public void init() {
        // 初始化指标计数器
        milvusFailureCounter = Counter.builder("service.failure")
                .tag("service", "milvus")
                .description("Milvus service failure count")
                .register(meterRegistry);

        milvusRecoveryCounter = Counter.builder("service.recovery")
                .tag("service", "milvus")
                .description("Milvus service recovery count")
                .register(meterRegistry);

        redisFailureCounter = Counter.builder("service.failure")
                .tag("service", "redis")
                .description("Redis service failure count")
                .register(meterRegistry);

        redisRecoveryCounter = Counter.builder("service.recovery")
                .tag("service", "redis")
                .description("Redis service recovery count")
                .register(meterRegistry);
    }

    /**
     * 定期检查服务健康状态
     * 每30秒执行一次
     */
    @Scheduled(fixedRate = 30000)
    public void checkServices() {
        // 检查Milvus服务健康状态
        checkMilvusHealth();
        
        // 检查Redis服务健康状态
        checkRedisHealth();
    }

    /**
     * 检查Milvus服务健康状态
     */
    private void checkMilvusHealth() {
        try {
            // 检查Milvus集合是否存在
            R<Boolean> response = milvusClient.hasCollection(
                    HasCollectionParam.newBuilder()
                            .withCollectionName(milvusConfig.getCollection())
                            .build()
            );
            
            // 检查响应状态
            boolean currentHealth = response.getStatus() == R.Status.Success.getCode();
            
            // 状态变化处理
            handleStateChange(milvusHealthy, currentHealth, "Milvus", 
                    milvusFailureCounter, milvusRecoveryCounter, this::reconnectMilvus);
            
        } catch (Exception e) {
            log.error("Milvus健康检查失败", e);
            if (milvusHealthy.compareAndSet(true, false)) {
                milvusFailureCounter.increment();
                log.warn("Milvus服务不可用，尝试恢复...");
                reconnectMilvus();
            }
        }
    }

    /**
     * 检查Redis服务健康状态
     */
    private void checkRedisHealth() {
        RedisConnection connection = null;
        try {
            // 获取Redis连接
            connection = redisConnectionFactory.getConnection();
            // 执行PING命令
            boolean currentHealth = "PONG".equals(connection.ping());
            
            // 状态变化处理
            handleStateChange(redisHealthy, currentHealth, "Redis", 
                    redisFailureCounter, redisRecoveryCounter, this::reconnectRedis);
            
        } catch (Exception e) {
            log.error("Redis健康检查失败", e);
            if (redisHealthy.compareAndSet(true, false)) {
                redisFailureCounter.increment();
                log.warn("Redis服务不可用，尝试恢复...");
                reconnectRedis();
            }
        } finally {
            // 关闭连接
            if (connection != null) {
                connection.close();
            }
        }
    }

    /**
     * 处理服务状态变化
     */
    private void handleStateChange(AtomicBoolean stateFlag, boolean currentHealth, 
                                  String serviceName, Counter failureCounter, 
                                  Counter recoveryCounter, Runnable recoveryAction) {
        if (currentHealth) {
            // 服务恢复
            if (stateFlag.compareAndSet(false, true)) {
                recoveryCounter.increment();
                log.info("{} 服务已恢复", serviceName);
            }
        } else {
            // 服务故障
            if (stateFlag.compareAndSet(true, false)) {
                failureCounter.increment();
                log.warn("{} 服务不可用，尝试恢复...", serviceName);
                recoveryAction.run();
            }
        }
    }

    /**
     * 重连Milvus服务
     */
    private void reconnectMilvus() {
        try {
            log.info("尝试重新连接Milvus服务...");
            // 实际实现可能需要根据具体部署方式调整
            // 例如：重新创建客户端、重新加载集合等
            
            // 这里简单地尝试检查集合是否存在
            R<Boolean> response = milvusClient.hasCollection(
                    HasCollectionParam.newBuilder()
                            .withCollectionName(milvusConfig.getCollection())
                            .build()
            );
            
            if (response.getStatus() == R.Status.Success.getCode()) {
                log.info("Milvus服务重连成功");
                milvusHealthy.set(true);
                milvusRecoveryCounter.increment();
            } else {
                log.warn("Milvus服务重连失败");
            }
        } catch (Exception e) {
            log.error("Milvus服务重连失败", e);
        }
    }

    /**
     * 重连Redis服务
     */
    private void reconnectRedis() {
        try {
            log.info("尝试重新连接Redis服务...");
            // 实际实现可能需要根据具体部署方式调整
            
            // 尝试获取新连接并执行PING命令
            RedisConnection connection = redisConnectionFactory.getConnection();
            try {
                if ("PONG".equals(connection.ping())) {
                    log.info("Redis服务重连成功");
                    redisHealthy.set(true);
                    redisRecoveryCounter.increment();
                } else {
                    log.warn("Redis服务重连失败");
                }
            } finally {
                connection.close();
            }
        } catch (Exception e) {
            log.error("Redis服务重连失败", e);
        }
    }
} 